package blockchain.block;

import java.util.ArrayList;
import java.util.List;

// 定义梅克尔树节点类
class MerkleTreeBlockNode {
    String hash;
    MerkleTreeBlockNode left;
    MerkleTreeBlockNode right;
    MerkleTreeBlockNode parent;

    Boolean isBlock = false;
    Block block;

    public MerkleTreeBlockNode(String value) {
        this.hash = Sha.Sha256(value);
    }
    public MerkleTreeBlockNode(Block block) {
        this.block = block;
        this.hash = block.getHash();
        this.isBlock = true;
    }
}

// 梅克尔树类
public class MerkleTreeBlock {
    public MerkleTreeBlockNode root;
    public List<Block> dataStr;
    public MerkleTreeBlock(List<Block> leafValues) {
        dataStr = new ArrayList<>();
        for (Block value : leafValues) {
            dataStr.add(value);
        }
        this.root = buildMerkleTree();
    }
    public MerkleTreeBlock() {
        dataStr = new ArrayList<>();
        dataStr.add(new Block("\"Genesis Block\""));
        this.root = buildMerkleTree();
    }

    // 构建初始梅克尔树
    private MerkleTreeBlockNode buildMerkleTree() {
        List<MerkleTreeBlockNode> leafNodes = new ArrayList<>();
        for (Block block : dataStr) {
            leafNodes.add(new MerkleTreeBlockNode(block));
        }

        List<MerkleTreeBlockNode> currentLevel = leafNodes;
        while (currentLevel.size() > 1) {
            List<MerkleTreeBlockNode> nextLevel = new ArrayList<>();
            for (int i = 0; i < currentLevel.size(); i += 2) {
                MerkleTreeBlockNode left = currentLevel.get(i);
                MerkleTreeBlockNode right = (i + 1 < currentLevel.size()) ? currentLevel.get(i + 1) : null;
                String parentValue;
                if (right != null) {
                    parentValue = left.hash + right.hash;
                } else {
                    parentValue = left.hash;
                }
                MerkleTreeBlockNode parent = new MerkleTreeBlockNode(parentValue);
                parent.left = left;
                parent.right = right;
                left.parent = parent;
                if (right != null) {
                    right.parent = parent;
                }
                nextLevel.add(parent);
            }
            currentLevel = nextLevel;
        }
        return currentLevel.isEmpty() ? null : currentLevel.get(0);
    }

    // 添加新的叶节点并更新梅克尔树
    public void addLeafNode(Block newLeafValue) {
        //List<String> leafValues = new ArrayList<>();
        //collectLeafValues(root, leafValues);
        //leafValues.add(newLeafValue);
        dataStr.add(newLeafValue);
        this.root = buildMerkleTree();
    }

    // 收集所有叶节点的值
    public void collectLeafValues(MerkleTreeBlockNode node, List<Block> leafValues) {
        if (node.left == null && node.right == null) {
            leafValues.add(node.block);
        }
        if (node.left != null) {
            collectLeafValues(node.left, leafValues);
        }
        if (node.right != null) {
            collectLeafValues(node.right, leafValues);
        }
    }

    // 获取根节点的值
    public String getRootValue() {
        return root != null ? root.hash : null;
    }

    public String print() {
        List<StringBuilder> builders = new ArrayList<>();
        builders.add(new StringBuilder());
        builders.get(0).append(root.hash);
        if (root.left != null) {
            DiGui(root.left, 1, builders);
        }
        if (root.right != null) {
            DiGui(root.right, 1, builders);
        }
        StringBuilder allBuilder = new StringBuilder();
        for (StringBuilder builder : builders) {
            allBuilder.append(builder).append("\n");
        }
        return allBuilder.toString();
    }

    public void DiGui(MerkleTreeBlockNode mTN, int depth, List<StringBuilder> builders) {
        if (depth == builders.size()) {
            builders.add(new StringBuilder());
        }
        builders.get(depth).append(mTN.hash).append("   ");
        if (mTN.left != null) {
            DiGui(mTN.left, depth + 1, builders);
        }
        if (mTN.right != null) {
            DiGui(mTN.right, depth + 1, builders);
        }
    }

    public static void main(String[] args) {
        List<Block> leafValues = List.of(new Block("A"), new Block("B"), new Block("C"), new Block("D"));
        MerkleTreeBlock merkleTreeBlock = new MerkleTreeBlock(leafValues);
        System.out.println("初始梅克尔树根节点值: " + merkleTreeBlock.getRootValue());

        System.out.println(merkleTreeBlock.print());

        Block newLeafValue = new Block("E");
        merkleTreeBlock.addLeafNode(newLeafValue);
        System.out.println("添加新节点后梅克尔树根节点值: " + merkleTreeBlock.getRootValue());

        System.out.println(merkleTreeBlock.print());
    }
}
